{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4e8d8398",
   "metadata": {},
   "source": [
    "# 02 Automatic Vectorization\n",
    "\n",
    "Original Documentation: https://docs.jax.dev/en/latest/automatic-vectorization.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "85aea56d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08e51446",
   "metadata": {},
   "source": [
    "## Manual vectorization\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "55d5d8f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[11. 20. 29.]\n"
     ]
    }
   ],
   "source": [
    "def convolve(x, w):\n",
    "    output = []\n",
    "    for i in range(1, len(x) - 1):\n",
    "        output.append(jnp.dot(x[i - 1 : i + 2], w))\n",
    "    return jnp.array(output)\n",
    "\n",
    "\n",
    "x = jnp.arange(5)\n",
    "w = jnp.array([2.0, 3.0, 4.0])\n",
    "\n",
    "print(convolve(x, w))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cf11ace",
   "metadata": {},
   "source": [
    "Suppose we wanted to apply `convolve` to a batch of weights `w` and vectors `x`. The naive implementation would introduce a new Python loop that calls `convolve`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "89f762dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[11. 20. 29.]\n",
      " [11. 20. 29.]]\n"
     ]
    }
   ],
   "source": [
    "def convolve(x, w):\n",
    "    output = []\n",
    "    for i in range(1, len(x) - 1):\n",
    "        output.append(jnp.dot(x[i - 1 : i + 2], w))\n",
    "    return jnp.array(output)\n",
    "\n",
    "\n",
    "x = jnp.arange(5)\n",
    "w = jnp.array([2.0, 3.0, 4.0])\n",
    "\n",
    "# Sample batch of vectors and weights\n",
    "xs = jnp.stack([x, x])\n",
    "ws = jnp.stack([w, w])\n",
    "\n",
    "\n",
    "def naive_batch_convolve(xs, ws):\n",
    "    assert xs.shape[0] == ws.shape[0]\n",
    "    output = []\n",
    "    for i in range(xs.shape[0]):\n",
    "        output.append(convolve(xs[i], ws[i]))\n",
    "    return jnp.stack(output)\n",
    "\n",
    "\n",
    "print(naive_batch_convolve(xs, ws))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d140205",
   "metadata": {},
   "source": [
    "This is quite inefficient since each batch’s convolution could have been done in parallel. To make it more efficient, we could rewrite to express the computations in a batched form:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "617a2872",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[11. 20. 29.]\n",
      " [11. 20. 29.]]\n"
     ]
    }
   ],
   "source": [
    "def manual_batch_convolve(xs, ws):\n",
    "    assert xs.shape[0] == ws.shape[0]\n",
    "    output = []\n",
    "    for i in range(1, xs.shape[1] - 1):  # xs.shape[-1] == 5\n",
    "        # Apply convolution along axis=1\n",
    "        # xs[:, i - 1 : i + 2] is sliding window of size 3 across axis=1\n",
    "        output.append(jnp.sum(xs[:, i - 1 : i + 2] * ws, axis=1))\n",
    "    return jnp.stack(output, axis=1)\n",
    "\n",
    "\n",
    "print(manual_batch_convolve(xs, ws))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96313640",
   "metadata": {},
   "source": [
    "Here, we compute the convolution for all windows in the batch simultaneously at each iteration by sliding a window along the last axis. In other words, each loop iteration processes the same window across all batches in parallel (since we use `jnp.sum()`)\n",
    "\n",
    "However, this is messy and error-prone.\n",
    "\n",
    "## Automatic vectorization\n",
    "\n",
    "`jax.vmap()` automatically generates a vectorized implementation of a function:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3b06b067",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[11. 20. 29.]\n",
      " [11. 20. 29.]]\n"
     ]
    }
   ],
   "source": [
    "auto_batch_convolve = jax.vmap(convolve)\n",
    "\n",
    "print(auto_batch_convolve(xs, ws))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "951d27a8",
   "metadata": {},
   "source": [
    "If the batch dimension is not the first axis, use the `in_axes` and `out_axes` arguments to specify the location of the batch dimension in inputs and outputs.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "dbc7d154",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[11. 11.]\n",
      " [20. 20.]\n",
      " [29. 29.]]\n"
     ]
    }
   ],
   "source": [
    "# Transpose so batch dimension is now axis=1\n",
    "xst = xs.transpose()  # shape (5, 2)\n",
    "wst = ws.transpose()  # shape (3, 2)\n",
    "\n",
    "auto_batch_convolve_v2 = jax.vmap(convolve, in_axes=1, out_axes=1)\n",
    "\n",
    "print(auto_batch_convolve_v2(xst, wst))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3013a8ba",
   "metadata": {},
   "source": [
    "## Combining transformations\n",
    "\n",
    "`jax.jit()` and `jax.vmap()` are composable; we can wrap them in either order.\n",
    "\n",
    "Notably, the order of wrapping them does matter: JIT compiling a vectorized function is different from vectorizing a JIT compiled function.\n",
    "\n",
    "Usually, you want to JIT compile a vectorized function (`jax.jit(jax.vmap(fn))`) since it gives the XLA compiler more code to optmize upfront.\n",
    "\n",
    "However, if the shapes of batches are not consistent, then we want to vectorize a JIT compiled function (`jax.vmap(jax.jit(fn))`) since we cannot provide different shapes to a function that has been compiled for one shape.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "5c86ac45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[11. 20. 29.]\n",
      " [11. 20. 29.]]\n"
     ]
    }
   ],
   "source": [
    "compiled_batch_convolve = jax.jit(\n",
    "    jax.vmap(convolve)\n",
    ")  # Shape along axis=1 is consistent\n",
    "\n",
    "print(compiled_batch_convolve(xs, ws))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
